package org.fujay.commons.langchain4j.rag.augmentor.retrieval;

import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import lombok.extern.slf4j.Slf4j;
import org.fujay.commons.langchain4j.rag.augmentor.EnhancedRetrieval;
import org.opensearch.client.opensearch.OpenSearchClient;
import org.opensearch.client.opensearch.core.SearchResponse;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import static dev.langchain4j.internal.ValidationUtils.ensureNotNull;

/**
 * 默认的增强检索器实现
 * 实现了EnhancedRetrieval接口，提供了以下功能：
 * 1. 混合检索：结合语义检索和关键词检索
 * 2. 多策略检索：支持BM25和TF-IDF算法
 * 3. 缓存机制：使用Guava Cache实现结果缓存--> 后续添加redis替代内存缓存
 * 4. 多维度评分：考虑时间衰减、长度归一化等因素
 * 5. 结果重排序：基于多维度评分进行结果重排序
 */
@Slf4j
public class DefaultEnhancedRetrieval implements EnhancedRetrieval {

    /**
     * 检索配置
     */
    private final Config config;
    /**
     * OpenSearch客户端
     */
    private final OpenSearchClient openSearchClient;
    /**
     * OpenSearch索引名称
     */
    private final String indexName;
    /**
     * 结果缓存
     */
    private final Cache<String, List<Content>> resultCache;
    /**
     * 缓存统计信息
     */
    private final CacheStats cacheStats;

    /**
     * 构造函数
     *
     * @param config           检索配置
     * @param openSearchClient OpenSearch客户端
     * @param indexName        OpenSearch索引名称
     */
    private DefaultEnhancedRetrieval(Config config, OpenSearchClient openSearchClient, String indexName) {
        this.config = ensureNotNull(config, "config");
        this.openSearchClient = ensureNotNull(openSearchClient, "openSearchClient");
        this.indexName = ensureNotNull(indexName, "indexName");

        // 初始化缓存
        this.resultCache = CacheBuilder.newBuilder()
                .maximumSize(config.getCacheSize())
                .expireAfterWrite(config.getCacheExpireTime(), TimeUnit.SECONDS)
                .recordStats()
                .build();

        this.cacheStats = CacheStats.builder()
                .hitCount(0)
                .missCount(0)
                .evictionCount(0)
                .size(0)
                .build();
    }

    /**
     * 创建构建器
     *
     * @return 构建器实例
     */
    public static Builder builder() {
        return new Builder();
    }

    /**
     * 执行检索
     *
     * @param query 查询对象
     * @return 检索结果列表
     */
    @Override
    public List<Content> retrieve(Query query) {
        // 1. 检查缓存
        String cacheKey = generateCacheKey(query);
        List<Content> cachedResults = resultCache.getIfPresent(cacheKey);
        if (cachedResults != null) {
            log.debug("从缓存中获取结果");
            cacheStats.setHitCount(cacheStats.getHitCount() + 1);
            return cachedResults;
        }
        cacheStats.setMissCount(cacheStats.getMissCount() + 1);

        List<Content> results = new ArrayList<>();

        // 2. 向量检索
        try {
            results.addAll(semanticSearch(query));
        } catch (Exception e) {
            log.error("语义检索失败", e);
        }

        // 3. 关键词检索
        if (config.isEnableHybridSearch()) {
            try {
                results.addAll(keywordSearch(query));
            } catch (Exception e) {
                log.error("关键词检索失败", e);
            }
        }

        // 4. 结果融合和重排序
        results = rerankResults(results, query);

        // 5. 缓存结果
        resultCache.put(cacheKey, results);
        cacheStats.setSize(resultCache.size());

        return results;
    }

    /**
     * 执行语义检索
     *
     * @param query 查询对象
     * @return 检索结果列表
     */
    private List<Content> semanticSearch(Query query) {
        // 1. 生成查询向量
        Embedding embeddedQuery = config.getEmbeddingModel().embed(query.text()).content();

        // 2. 构建检索请求
        EmbeddingSearchRequest searchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(embeddedQuery)
                .maxResults(config.getMaxResultsProvider().apply(query))
                .minScore(config.getMinScoreProvider().apply(query))
                .filter(config.getFilterProvider().apply(query))
                .build();

        // 3. 执行向量检索
        EmbeddingSearchResult<TextSegment> searchResult = config.getEmbeddingStore().search(searchRequest);

        // 4. 转换结果
        return searchResult.matches().stream()
                .map(embeddingMatch -> {
                    TextSegment segment = embeddingMatch.embedded();
                    Map<String, Object> metadataMap = new HashMap<>(segment.metadata().toMap());
                    metadataMap.put("score", String.valueOf(embeddingMatch.score() * config.getSemanticWeight()));
                    metadataMap.put("embeddingId", embeddingMatch.embeddingId());
                    metadataMap.put("source", "semantic_search");
                    TextSegment newSegment = TextSegment.from(segment.text(), Metadata.from(metadataMap));
                    return Content.from(newSegment);
                })
                .collect(Collectors.toList());
    }

    /**
     * 执行关键词检索
     *
     * @param query 查询对象
     * @return 检索结果列表
     */
    private List<Content> keywordSearch(Query query) throws IOException {
        // 1. 构建查询
        SearchResponse<Map> response = openSearchClient.search(builder -> builder
                        .index(indexName)
                        .query(q -> q
                                .multiMatch(m -> m
                                        .query(query.text())
                                        .fields("content^1.0")))
                        .size(config.getMaxResultsProvider().apply(query)),
                Map.class);

        // 2. 转换结果
        return response.hits().hits().stream()
                .map(hit -> {
                    @SuppressWarnings("unchecked")
                    Map<String, Object> source = (Map<String, Object>) hit.source();
                    Map<String, Object> metadataMap = Map.of(
                            "score", String.valueOf(hit.score() * config.getKeywordWeight()),
                            "source", "keyword_search");
                    return Content.from(TextSegment.from(
                            source.get("content").toString(),
                            Metadata.from(metadataMap)));
                })
                .collect(Collectors.toList());
    }

    /**
     * 生成缓存键
     *
     * @param query 查询对象
     * @return 缓存键
     */
    private String generateCacheKey(Query query) {
        return query.text() + "_" + query.metadata().toString();
    }

    /**
     * 重排序检索结果
     *
     * @param results 原始结果列表
     * @param query   查询对象
     * @return 重排序后的结果列表
     */
    private List<Content> rerankResults(List<Content> results, Query query) {
        // 1. 去重
        results = results.stream()
                .distinct()
                .collect(Collectors.toList());

        // 2. 多维度评分
        results = results.stream()
                .map(content -> {
                    double finalScore = calculateFinalScore(content, query);
                    Map<String, Object> metadataMap = new HashMap<>(content.textSegment().metadata().toMap());
                    metadataMap.put("score", String.valueOf(finalScore));
                    TextSegment segment = TextSegment.from(content.textSegment().text(),
                            Metadata.from(metadataMap));
                    return Content.from(segment);
                })
                .collect(Collectors.toList());

        // 3. 重排序
        return results.stream()
                .sorted(Comparator.comparingDouble(
                                content -> ((Content) content).textSegment().metadata().getDouble("score"))
                        .reversed())
                .limit(config.getMaxResultsProvider().apply(query))
                .collect(Collectors.toList());
    }

    /**
     * 计算最终得分
     *
     * @param content 内容对象
     * @param query   查询对象
     * @return 最终得分
     */
    private double calculateFinalScore(Content content, Query query) {
        // 1. 基础得分
        double baseScore = content.textSegment().metadata().getDouble("score");

        // 2. 来源权重
        double sourceWeight = Objects.equals(content.textSegment().metadata().getString("source"), "semantic_search")
                ? config.getSemanticWeight()
                : config.getKeywordWeight();

        // 3. 时间衰减因子
        Map<String, Object> metadataMap = content.textSegment().metadata().toMap();
        String timestamp = (String) metadataMap.getOrDefault("timestamp", String.valueOf(System.currentTimeMillis()));
        double timeDecay = Math
                .exp(-0.1 * (System.currentTimeMillis() - Long.parseLong(timestamp)) / (24 * 60 * 60 * 1000.0));

        // 4. 长度归一化因子
        double lengthNormalization = 1.0 / Math.sqrt(content.textSegment().text().length());

        // 5. 计算最终得分
        return baseScore * sourceWeight * timeDecay * lengthNormalization;
    }

    /**
     * 预热缓存
     *
     * @param queries 预热查询列表
     */
    @Override
    public void warmupCache(List<Query> queries) {
        queries.forEach(this::retrieve);
    }

    /**
     * 清除缓存
     */
    @Override
    public void clearCache() {
        resultCache.invalidateAll();
        cacheStats.setEvictionCount(cacheStats.getEvictionCount() + 1);
    }

    /**
     * 获取缓存统计信息
     *
     * @return 缓存统计信息
     */
    @Override
    public CacheStats getCacheStats() {
        return cacheStats;
    }

    /**
     * 构建器类
     */
    public static class Builder {
        private Config config;
        private OpenSearchClient openSearchClient;
        private String indexName;

        /**
         * 设置配置
         *
         * @param config 检索配置
         * @return 构建器实例
         */
        public Builder config(Config config) {
            this.config = config;
            return this;
        }

        /**
         * 设置OpenSearch客户端
         *
         * @param openSearchClient OpenSearch客户端
         * @return 构建器实例
         */
        public Builder openSearchClient(OpenSearchClient openSearchClient) {
            this.openSearchClient = openSearchClient;
            return this;
        }

        /**
         * 设置OpenSearch索引名称
         *
         * @param indexName OpenSearch索引名称
         * @return 构建器实例
         */
        public Builder indexName(String indexName) {
            this.indexName = indexName;
            return this;
        }

        /**
         * 构建检索器实例
         *
         * @return 检索器实例
         */
        public DefaultEnhancedRetrieval build() {
            return new DefaultEnhancedRetrieval(
                    config,
                    openSearchClient,
                    indexName);
        }
    }
}